import torch
import torch.nn as nn


# 定义一个基于LSTM的情绪分类模型
class EmotionLSTM(nn.Module):
    # 初始化模型的结构
    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
        super(EmotionLSTM, self).__init__()

        # LSTM 层：负责处理时序数据。参数包括输入维度、隐藏层维度、LSTM层数
        # batch_first=True 表示输入和输出的张量形状为 (batch_size, sequence_length, input_dim)
        # 在 LSTM 层之间添加 dropout 层，防止过拟合，丢弃 20% 的连接。
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True, dropout=0.2)

        # 全连接层：将LSTM的输出映射到情绪分类的输出空间，output_dim是情绪类别的数量（8个）
        self.fc = nn.Linear(hidden_dim, output_dim)

        # Softmax 层：将LSTM的输出转化为概率分布，输出每个情绪类别的概率
        self.softmax = nn.Softmax(dim=1)

    # 前向传播函数：定义数据流向模型的过程
    def forward(self, x):
        # x的形状是 (batch_size, sequence_length, input_dim)

        # 通过LSTM层进行计算，lstm_out包含了所有时间步的输出，h_n和c_n是LSTM的最后隐藏状态和记忆状态
        lstm_out, (h_n, c_n) = self.lstm(x)  # lstm_out形状为 (batch_size, sequence_length, hidden_dim)

        # LSTM输出的是每个时间步的隐藏状态，我们通常只关心最后一个时间步的输出
        last_hidden_state = lstm_out[:, -1, :]  # 取最后一个时间步的输出，形状是 (batch_size, hidden_dim)

        # 将LSTM的最后输出传入全连接层，得到最终的分类结果
        out = self.fc(last_hidden_state)  # 形状是 (batch_size, output_dim)

        # 通过Softmax层将输出转化为概率分布，输出每个情绪类别的概率
        return self.softmax(out)  # 形状是 (batch_size, output_dim)
